import torch
import torch.nn as nn
import torch.nn.functional as F


class AudioCNN(nn.Module):
    def __init__(self, num_classes, in_channels=1, feature_dim=64):
        """
        音频CNN模型
        参数:
            num_classes (int): 输出类别数
            in_channels (int): 输入通道数 (1 for mono, 3 for STFT)
            feature_dim (int): 特征维度 (根据输入频谱图大小调整)
        """
        super(AudioCNN, self).__init__()

        # 特征提取层
        self.features = nn.Sequential(
            # 输入形状: (batch, in_channels, freq_bins, time_steps)
            nn.Conv2d(in_channels, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))  # 全局平均池化
        )

        # 分类器
        self.classifier = nn.Sequential(
            nn.Linear(256, feature_dim),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(feature_dim, num_classes)
        )

    def forward(self, x):
        """输入形状: (batch, channels, freq_bins, time_steps)"""
        x = self.features(x)  # -> (batch, 256, 1, 1)
        x = torch.flatten(x, 1)  # -> (batch, 256)
        x = self.classifier(x)  # -> (batch, num_classes)
        return x


# 辅助函数：快速测试模型
def test_model():
    # 模拟输入 (batch=2, 1 channel, 128 freq bins, 100 time steps)
    x = torch.randn(2, 1, 128, 100)
    model = AudioCNN(num_classes=10)
    out = model(x)
    print("Output shape:", out.shape)  # 应该得到 [2, 10]


if __name__ == "__main__":
    test_model()